# coding = utf-8

'''
从3DIRCADB数据集中分别读取Dice
'''

import os
import SimpleITK as sitk
from PIL import Image
import numpy as np


def find_data_and_compute_dice(case_id):
    root_path = "F:\Dataset\Liver\\3Dircadb"
    root_path = os.path.join(root_path, "case_{}".format(str(case_id).zfill(5)))

    liver_and_tumor_path = os.path.join(root_path, "segmentation")
    raw_path = os.path.join(root_path, "imaging")

    ours_root = "F:\predict\\3dircadb\ours"
    hdenseunet_root = "F:\predict\\3dircadb\hdenseunet"
    munet_root = "F:\predict\\3dircadb\munet"
    unet_root = "F:\predict\\3dircadb\\unet"

    ours_root = os.path.join(ours_root, "case_{}\predict_tumor".format(str(case_id).zfill(5)))
    hdenseunet_root = os.path.join(hdenseunet_root, "case_{}\predict_tumor".format(str(case_id).zfill(5)))
    munet_root = os.path.join(munet_root, "case_{}\predict_tumor".format(str(case_id).zfill(5)))
    unet_root = os.path.join(unet_root, "case_{}\predict_tumor".format(str(case_id).zfill(5)))

    for i in range(len(os.listdir(liver_and_tumor_path))):
        liver_and_tumor = os.path.join(liver_and_tumor_path, "{}.npy".format(str(i).zfill(3)))
        liver_and_tumor = np.load(liver_and_tumor)
        liver = np.zeros(liver_and_tumor.shape)
        liver[liver_and_tumor > 0] = 1
        tumor = np.zeros(liver_and_tumor.shape)
        tumor[liver_and_tumor == 2] = 1

        raw_data = os.path.join(raw_path, "{}.npy".format(str(i).zfill(3)))
        raw_data = np.load(raw_data)

        if tumor.sum() <= 0:
            continue

        ours_image = os.path.join(ours_root, "{}.png".format(str(i).zfill(3)))
        hdenseunet_image = os.path.join(hdenseunet_root, "{}.png".format(str(i).zfill(3)))
        munet_image = os.path.join(munet_root, "{}.png".format(str(i).zfill(3)))
        unet_image = os.path.join(unet_root, "{}.png".format(str(i).zfill(3)))

        if os.path.exists(ours_image):
            ours = Image.open(ours_image).convert("L")
            ours = np.array(ours)
            ours[ours > 0] = 1
        else:
            ours = np.zeros(tumor.shape)
        dice_ours = float(2 * (ours * tumor).sum()) / float(ours.sum() + tumor.sum())

        if os.path.exists(hdenseunet_image):
            hdenseunet = Image.open(hdenseunet_image).convert("L")
            hdenseunet = np.array(hdenseunet)
            hdenseunet[hdenseunet > 0] = 1
        else:
            hdenseunet = np.zeros(tumor.shape)
        dice_hdenseunet = float(2 * (hdenseunet * tumor).sum()) / float(hdenseunet.sum() + tumor.sum())

        if os.path.exists(munet_image):
            munet = Image.open(munet_image).convert("L")
            munet = np.array(munet)
            munet[munet > 0] = 1
        else:
            munet = np.zeros(tumor.shape)
        dice_munet = float(2 * (munet * tumor).sum()) / float(munet.sum() + tumor.sum())

        if os.path.exists(unet_image):
            unet = Image.open(unet_image).convert("L")
            unet = np.array(unet)
            unet[unet > 0] = 1
        else:
            unet = np.zeros(tumor.shape)
        dice_unet = float(2 * (unet * tumor).sum()) / float(unet.sum() + tumor.sum())

        if dice_ours > dice_hdenseunet and dice_ours > dice_munet and dice_ours > dice_unet:
              print(i, round(dice_ours, 2), round(dice_hdenseunet, 2), round(dice_munet, 2), round(dice_unet, 2))









if __name__ == '__main__':
    find_data_and_compute_dice(case_id=18)